#include <stdio.h>
#include "picture_codec.h"

 
 
int flush_encoder(AVFormatContext *fmt_ctx,unsigned int stream_index)
{
    int ret;
    int got_frame;
	AVPacket enc_pkt;
    if (!(fmt_ctx->streams[stream_index]->codec->codec->capabilities & AV_CODEC_CAP_DELAY))
        return 0;
    while (1) {
        av_log(NULL, AV_LOG_INFO, "Flushing stream #%u encoder\n", stream_index);
        //ret = encode_write_frame(NULL, stream_index, &got_frame);
        enc_pkt.data = NULL;
		enc_pkt.size = 0;
		av_init_packet(&enc_pkt);
		ret = avcodec_encode_video2 (fmt_ctx->streams[stream_index]->codec, &enc_pkt, NULL, &got_frame);
		av_frame_free(NULL);
		if (ret < 0)
			break;
		if (!got_frame)
		{
			ret=0;
			break;
		}		
		
		av_log(NULL, AV_LOG_DEBUG, "Muxing frame\n");
		/* mux encoded frame */
		ret = av_write_frame(fmt_ctx, &enc_pkt);
		if (ret < 0)
            break;
    }
    return ret;
}
 


int picture_codec_process(codec_context_t  *codec)
{
	int ret;
	unsigned int stream_index;
	int got_frame,enc_got_frame;
	codec_io_context_t  *in_context  = &codec->in_context;
	codec_io_context_t  *out_context = &codec->out_context;
	picture_filter_t *filter               = &codec->filter;
	AVPacket *inPacket;
	AVPacket outPacket;

	// 包初始化 
	inPacket = (AVPacket *)av_malloc(sizeof(AVPacket));
	int picture_size = avpicture_get_size(in_context->pix_fmt, in_context->width, in_context->height);
	av_new_packet(&outPacket, picture_size);

	/* read all packets */
	while (1) {
		/* 读取包 */
		if ((ret = av_read_frame(in_context->fmt_ctx, inPacket)) < 0)
			break;
		stream_index = inPacket->stream_index;
		if(0 != stream_index)
			continue;

		/* 包解码 */
		ret = avcodec_decode_video2(in_context->fmt_ctx->streams[stream_index]->codec, 
								in_context->frame, 
								&got_frame, inPacket);
		if (ret < 0) {
			av_log(NULL, AV_LOG_ERROR, "Decoding failed\n");
			break;
		}	

		if (! got_frame) 
		{
			return -1;
		}

		/* 帧过滤 */
		picture_filter(filter, in_context->frame, out_context->frame);

		/* 包编码 */
		ret = avcodec_encode_video2 (out_context->fmt_ctx->streams[stream_index]->codec, 
								&outPacket, 
								out_context->frame, 
								&enc_got_frame);
		
		if (!enc_got_frame)
			continue;
	
		/* 写入报文 */
		av_write_frame(out_context->fmt_ctx, &outPacket);
		if (ret < 0)
		{
			av_log(NULL, AV_LOG_INFO, "Error write frame failed \n");
			return -1;
		}
		av_free_packet(inPacket);
		av_free_packet(&outPacket);
	}
 
	/* flush encoders */
	for (int i = 0; i < 1; i++) {
		/* flush encoder */
		ret = flush_encoder(out_context->fmt_ctx, i);
		if (ret < 0) {
			av_log(NULL, AV_LOG_ERROR, "Flushing encoder failed\n");
			return -1;
		}
	}
	av_write_trailer(out_context->fmt_ctx);
	return 0;
}





int picture_codec_input_init(codec_io_context_t *in_context, void *args, codec_callback ReadFun)
{
    int ret;
	in_context->buffer = (unsigned char*)av_malloc(32768);
	in_context->avio =avio_alloc_context(in_context->buffer, 32768,0, args, ReadFun, NULL, NULL);  
	if(in_context->avio==NULL)
	{
		return -1;
	}

	in_context->fmt_ctx = avformat_alloc_context();
	if(NULL == in_context->fmt_ctx)
	{
		return -1;
	}
	in_context->fmt_ctx->pb = in_context->avio; 
	in_context->fmt_ctx->flags = AVFMT_FLAG_CUSTOM_IO;
	if ((ret = avformat_open_input(&in_context->fmt_ctx, "whatever", NULL, NULL)) < 0) {
		av_log(NULL, AV_LOG_ERROR, "Cannot open input file\n");
		return -1;
	}

	if ((ret = avformat_find_stream_info(in_context->fmt_ctx, NULL)) < 0) {
		av_log(NULL, AV_LOG_ERROR, "Cannot find stream information\n");
		av_log(NULL, AV_LOG_ERROR, "Cannot open input file\n");
		return ret;
	}
	for (int i = 0; i < in_context->fmt_ctx->nb_streams; i++) {
		AVStream *stream;
		AVCodecContext *codec_ctx;
		stream = in_context->fmt_ctx->streams[i];
		codec_ctx = stream->codec;
		/* Reencode video & audio and remux subtitles etc. */
		if (codec_ctx->codec_type == AVMEDIA_TYPE_VIDEO){
			/* Open decoder */
			ret = avcodec_open2(codec_ctx, avcodec_find_decoder(codec_ctx->codec_id), NULL);
			if (ret < 0) {
				av_log(NULL, AV_LOG_ERROR, "Failed to open decoder for stream #%u\n", i);
				return ret;
			}
		}
		in_context->channels       = codec_ctx->channels;
		in_context->channel_layout = codec_ctx->channel_layout;
		in_context->width          = codec_ctx->width;
		in_context->height         = codec_ctx->height;
		in_context->pix_fmt        = codec_ctx->pix_fmt;
	}

	in_context->frame = av_frame_alloc();
	if(NULL == in_context->frame)
	{
		av_log(NULL, AV_LOG_ERROR, "Error alloc frame failed \n");
		return -1;
	}
	return 0;	
}



int picture_codec_output_init(codec_io_context_t *output_context, void *args, codec_callback WriteFun, AVFormatContext *iformat_ctx)
{
    int ret;
	AVCodec  *encoder;
	AVStream *out_stream;
	AVStream *in_stream;
	AVCodecContext *dec_ctx, *enc_ctx;
		
	avformat_alloc_output_context2(&output_context->fmt_ctx, NULL, "image2pipe", NULL);
 	
	output_context->buffer=(unsigned char*)av_malloc(32768);
	output_context->avio =avio_alloc_context(output_context->buffer, 32768, 1, args, NULL,WriteFun,NULL);  
	if(output_context->avio==NULL)
	{
		return -1;
	}

	//avio_out->write_packet=write_packet;
	output_context->fmt_ctx->pb = output_context->avio; 
	output_context->fmt_ctx->flags=AVFMT_FLAG_CUSTOM_IO;
	for (int i = 0; i < 1; i++) {
		out_stream = avformat_new_stream(output_context->fmt_ctx, NULL);
		if (!out_stream) {
			av_log(NULL, AV_LOG_ERROR, "Failed allocating output stream\n");
			return AVERROR_UNKNOWN;
		}
		in_stream = iformat_ctx->streams[i];
		dec_ctx = in_stream->codec;
		enc_ctx = out_stream->codec;
		if (dec_ctx->codec_type == AVMEDIA_TYPE_VIDEO)
		{
			encoder = avcodec_find_encoder(AV_CODEC_ID_MJPEG);
			enc_ctx->height = dec_ctx->height;
			enc_ctx->width = dec_ctx->width;
			enc_ctx->sample_aspect_ratio = dec_ctx->sample_aspect_ratio;
			enc_ctx->pix_fmt = encoder->pix_fmts[0];
			enc_ctx->time_base = dec_ctx->time_base;
			//enc_ctx->time_base.num = 1;
			//enc_ctx->time_base.den = 25;
			//H264的必备选项，没有就会错
			enc_ctx->me_range=16;
			enc_ctx->max_qdiff = 4;
			enc_ctx->qmin = 10;
			enc_ctx->qmax = 51;
			enc_ctx->qcompress = 0.6; 
			enc_ctx->refs=3;
			enc_ctx->bit_rate = 500000;
 
			ret = avcodec_open2(enc_ctx, encoder, NULL);
			if (ret < 0) {
				av_log(NULL, AV_LOG_ERROR, "Cannot open video encoder for stream #%u\n", i);
				return ret;
			}
		}
		else if (dec_ctx->codec_type == AVMEDIA_TYPE_UNKNOWN) {
			av_log(NULL, AV_LOG_FATAL, "Elementary stream #%d is of unknown type, cannot proceed\n", i);
			return AVERROR_INVALIDDATA;
		} else {
			/* if this stream must be remuxed */
			ret = avcodec_copy_context(output_context->fmt_ctx->streams[i]->codec, iformat_ctx->streams[i]->codec);
			if (ret < 0) {
				av_log(NULL, AV_LOG_ERROR, "Copying stream context failed\n");
				return ret;
			}
		}
		if (output_context->fmt_ctx->oformat->flags & AVFMT_GLOBALHEADER)
			enc_ctx->flags |= AV_CODEC_FLAG_GLOBAL_HEADER;
	}
	
	/* init muxer, write output file header */
	ret = avformat_write_header(output_context->fmt_ctx, NULL);
	if (ret < 0) {
		av_log(NULL, AV_LOG_ERROR, "Error occurred when opening output file\n");
		return ret;
	}

	output_context->frame = av_frame_alloc();
	if(NULL == output_context->frame)
	{
		av_log(NULL, AV_LOG_ERROR, "Error alloc frame failed \n");
		return -1;
	}
#if 0
	output_context->packet = (AVPacket *)av_malloc(sizeof(AVPacket));
	if(NULL == output_context->packet)
	{
		av_log(NULL, AV_LOG_ERROR, "Error alloc packet failed \n");
		return -1;
	}
#endif	
	return 0;
}







int picture_codec_init(codec_context_t  *codec, void *args, codec_callback ReadFun, codec_callback WriteFun, const char *filter_descr)
{
	av_register_all();
	memset(codec, 0, sizeof(codec_context_t));

	if(0 != picture_codec_input_init(&codec->in_context, args, ReadFun))
	{
		picture_codec_exit(codec);
		return -1;
	}

	if(0 != picture_codec_output_init(&codec->out_context, args, WriteFun, codec->in_context.fmt_ctx))
	{
		picture_codec_exit(codec);
		return -1;
	}

	picture_filter_init(&codec->filter, filter_descr, 
					codec->in_context.width, 
					codec->in_context.height, 
					codec->in_context.pix_fmt);
	
    return 0;
}



int picture_codec_exit(codec_context_t  *codec)
{	
	;
	if(codec->in_context.avio)
		av_freep(codec->in_context.avio);
#if 0			
	if(codec->in_context.packet)
		av_free_packet(codec->in_context.packet);
#endif		
	if(codec->in_context.frame)
		av_frame_free(&codec->in_context.frame);
	if(codec->in_context.fmt_ctx)
	{
		avformat_close_input(&codec->in_context.fmt_ctx);
		codec->in_context.buffer = NULL;
	}		
	if(codec->in_context.buffer)
		av_free(codec->in_context.buffer);


	if(codec->out_context.avio)
		av_freep(codec->out_context.avio);
	if(codec->out_context.buffer)
		av_free(codec->out_context.buffer);
#if 0		
	if(codec->out_context.packet)
		av_free_packet(codec->out_context.packet);	
#endif		
	if(codec->out_context.frame)
		av_frame_free(&codec->out_context.frame);
	if(codec->out_context.fmt_ctx)
		avformat_free_context(codec->out_context.fmt_ctx);
	

	picture_filter_exit(&codec->filter);

	memset(codec, 0, sizeof(codec_context_t));
	// fcloseall();
    return 0;
}

